import logging
import sys
from typing import Any, List

import orjson
from hamilton import base
from hamilton.async_driver import AsyncDriver
from haystack import component
from haystack.components.builders.prompt_builder import PromptBuilder
from langfuse.decorators import observe
from pydantic import BaseModel

from src.core.pipeline import BasicPipeline
from src.core.provider import LLMProvider

logger = logging.getLogger("wren-ai-service")


sql_summary_system_prompt = """
### TASK ###
You are a great data analyst. You are now given a task to summarize a list SQL queries in a human-readable format where each summary should be within 10-20 words.
You will be given a list of SQL queries and a user's question.

### INSTRUCTIONS ###
- SQL query summary must be within 10-20 words.
- SQL query summary must be human-readable and easy to understand.
- SQL query summary must be concise and to the point.
- SQL query summary must be in the same language user specified.

### OUTPUT FORMAT ###
Please return the result in the following JSON format:

{
    "sql_summary_results": [
        {
            "summary": <SQL_QUERY_SUMMARY_STRING>
        }
    ]
}
"""

sql_summary_user_prompt_template = """
User's Question: {{query}}
SQLs: {{sqls}}
Language: {{language}}

Please think step by step.
"""


@component
class SQLSummaryPostProcessor:
    @component.output_types(
        sql_summary_results=List[str],
    )
    def run(self, sqls: List[str], replies: List[str]):
        try:
            # 检查replies是否有效
            if not replies or not isinstance(replies, list) or len(replies) == 0:
                logger.error("SQL Summary Post Processor: Empty or invalid replies")
                return {
                    "sql_summary_results": [],
                }
            
            # 获取第一个回复
            reply = replies[0]
            if not isinstance(reply, str) or reply.strip() == "":
                logger.error("SQL Summary Post Processor: Empty reply content")
                return {
                    "sql_summary_results": [],
                }
            
            logger.debug(f"Original reply: {reply[:200]}...")
            
            # 尝试多种方式解析JSON
            summary_results = None
            
            # 尝试直接解析
            try:
                parsed_reply = orjson.loads(reply)
                logger.debug(f"Successfully parsed reply as JSON: {str(parsed_reply)[:200]}...")
                summary_results = parsed_reply.get("sql_summary_results", [])
            except orjson.JSONDecodeError as e:
                logger.warning(f"Initial JSON parse failed: {e}. Trying to fix...")
                
                # 尝试找到JSON部分
                try:
                    # 查找可能的JSON开始位置 (通常是 '{')
                    json_start = reply.find('{')
                    json_end = reply.rfind('}')
                    
                    if json_start >= 0 and json_end > json_start:
                        json_content = reply[json_start:json_end + 1]
                        logger.info(f"Trying to parse extracted JSON: {json_content[:100]}...")
                        parsed_reply = orjson.loads(json_content)
                        summary_results = parsed_reply.get("sql_summary_results", [])
                        logger.info("Successfully parsed extracted JSON")
                    else:
                        # 尝试使用正则表达式匹配所有JSON对象
                        import re
                        json_pattern = r'\{(?:[^{}]|(?:\{(?:[^{}]|(?:\{[^{}]*\}))*\}))*\}'
                        matches = re.findall(json_pattern, reply)
                        
                        if matches:
                            # 尝试解析每一个匹配的JSON
                            for json_candidate in matches:
                                try:
                                    parsed_json = orjson.loads(json_candidate)
                                    if "sql_summary_results" in parsed_json:
                                        summary_results = parsed_json["sql_summary_results"]
                                        logger.info("Successfully parsed JSON from regex pattern")
                                        break
                                except orjson.JSONDecodeError:
                                    continue
                except Exception as nested_e:
                    logger.error(f"Failed to extract JSON: {nested_e}")
            
            # 如果仍然无法解析，尝试从文本中提取摘要
            if summary_results is None:
                logger.warning("Could not parse JSON properly, attempting to extract summaries from text")
                
                # 尝试寻找摘要模式，例如 "summary": "some text"
                import re
                summary_pattern = r'"summary"\s*:\s*"([^"]+)"'
                matches = re.findall(summary_pattern, reply)
                
                if matches:
                    logger.info(f"Found {len(matches)} summaries using regex")
                    summary_results = [{"summary": match} for match in matches]
                else:
                    # 如果仍然找不到，创建一个简单的摘要
                    logger.warning("Could not extract summaries, creating generic summaries")
                    summary_results = []
            
            # 如果SQL比摘要多，为剩余的SQL创建通用摘要
            if len(sqls) > len(summary_results):
                logger.warning(f"Only {len(summary_results)} summaries for {len(sqls)} SQLs, filling in generic summaries")
                for _ in range(len(sqls) - len(summary_results)):
                    summary_results.append({"summary": "SQL操作"})
            
            # 确保每个摘要都有summary字段
            for i, result in enumerate(summary_results):
                if not isinstance(result, dict) or "summary" not in result:
                    summary_results[i] = {"summary": "SQL操作"}
            
            # 创建最终结果
            return {
                "sql_summary_results": [
                    {"sql": sql, "summary": summary.get("summary", "SQL操作")}
                    for (sql, summary) in zip(sqls, summary_results[:len(sqls)])
                ],
            }
            
        except Exception as e:
            logger.exception(f"Error in SQLSummaryPostProcessor: {e}")
            
            # 生成默认摘要
            return {
                "sql_summary_results": [
                    {"sql": sql, "summary": f"SQL操作 {i+1}"}
                    for i, sql in enumerate(sqls)
                ],
            }


## Start of Pipeline
@observe(capture_input=False)
def prompt(
    query: str,
    sqls: List[str],
    language: str,
    prompt_builder: PromptBuilder,
) -> dict:
    return prompt_builder.run(
        query=query,
        sqls=sqls,
        language=language,
    )


@observe(as_type="generation", capture_input=False)
async def generate_sql_summary(prompt: dict, generator: Any) -> dict:
    return await generator(prompt=prompt.get("prompt"))


def post_process(
    generate_sql_summary: dict,
    sqls: List[str],
    post_processor: SQLSummaryPostProcessor,
) -> dict:
    return post_processor.run(sqls, generate_sql_summary.get("replies"))


## End of Pipeline
class SummaryResult(BaseModel):
    summary: str


class SummaryResults(BaseModel):
    sql_summary_results: list[SummaryResult]


SQL_SUMMARY_MODEL_KWARGS = {
    "response_format": {
        "type": "json_schema",
        "json_schema": {
            "name": "sql_summary",
            "schema": SummaryResults.model_json_schema(),
        },
    }
}


class SQLSummary(BasicPipeline):
    def __init__(
        self,
        llm_provider: LLMProvider,
        **kwargs,
    ):
        self._components = {
            "generator": llm_provider.get_generator(
                system_prompt=sql_summary_system_prompt,
                generation_kwargs=SQL_SUMMARY_MODEL_KWARGS,
            ),
            "prompt_builder": PromptBuilder(template=sql_summary_user_prompt_template),
            "post_processor": SQLSummaryPostProcessor(),
        }

        super().__init__(
            AsyncDriver({}, sys.modules[__name__], result_builder=base.DictResult())
        )

    @observe(name="SQL Summary")
    async def run(
        self,
        query: str,
        sqls: List[str],
        language: str,
    ):
        logger.info("SQL Summary pipeline is running...")
        return await self._pipe.execute(
            ["post_process"],
            inputs={
                "query": query,
                "sqls": sqls,
                "language": language,
                **self._components,
            },
        )


if __name__ == "__main__":
    from src.pipelines.common import dry_run_pipeline

    dry_run_pipeline(
        SQLSummary,
        "sql_summary",
        query="this is a test query",
        sqls=[],
        language="English",
    )
